"""
后处理可视化模块

本模块负责计算结果的可视化，包括：
1. 子通道分布图
2. 轴向分布图
3. 时间历程图
4. 导出数据

作者：[您的名字]
日期：[创建日期]
"""

import os
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List
from ..physics.two_fluid_model import FluidState
from ..utils.logger import get_logger

class ResultVisualizer:
    """计算结果可视化类"""
    
    def __init__(self, config: dict):
        """
        初始化可视化器
        
        参数:
            config: dict, 配置字典
        """
        self.logger = get_logger()
        self.config = config
        
        # 创建输出目录
        self.output_dir = os.path.join(
            config['output']['directory'],
            'figures'
        )
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 设置绘图样式
        plt.style.use('seaborn')
        
    def plot_results(self, state: Dict[str, FluidState], step: int, time: float):
        """
        绘制计算结果
        
        参数:
            state: Dict[str, FluidState], 当前状态
            step: int, 当前步数
            time: float, 当前时间
        """
        # 绘制子通道分布图
        self._plot_subchannel_distribution(state, step, time)
        
        # 绘制轴向分布图
        self._plot_axial_distribution(state, step, time)
        
        # 绘制时间历程
        self._plot_time_history(state, step, time)
        
    def _plot_subchannel_distribution(
        self,
        state: Dict[str, FluidState],
        step: int,
        time: float
    ):
        """
        绘制子通道分布图
        
        参数:
            state: Dict[str, FluidState], 当前状态
            step: int, 当前步数
            time: float, 当前时间
        """
        # 创建图形
        fig, axes = plt.subplots(2, 2, figsize=(12, 12))
        fig.suptitle(f'子通道分布 (t = {time:.3f} s)', fontsize=14)
        
        # 获取通道布局
        layout = self.config['boundary_conditions']['subchannel_layout']
        n_rows = len(layout)
        n_cols = len(layout[0])
        
        # 准备数据
        void_fraction = state['vapor'].alpha.reshape(n_rows, n_cols)
        pressure = state['pressure'].reshape(n_rows, n_cols)
        liquid_temp = state['liquid'].T.reshape(n_rows, n_cols)
        vapor_temp = state['vapor'].T.reshape(n_rows, n_cols)
        
        # 绘制空泡率分布
        im0 = axes[0,0].imshow(void_fraction, cmap='coolwarm')
        axes[0,0].set_title('空泡率分布')
        plt.colorbar(im0, ax=axes[0,0])
        
        # 绘制压力分布
        im1 = axes[0,1].imshow(pressure/1e6, cmap='viridis')  # MPa
        axes[0,1].set_title('压力分布 (MPa)')
        plt.colorbar(im1, ax=axes[0,1])
        
        # 绘制液相温度分布
        im2 = axes[1,0].imshow(liquid_temp-273.15, cmap='coolwarm')  # °C
        axes[1,0].set_title('液相温度分布 (°C)')
        plt.colorbar(im2, ax=axes[1,0])
        
        # 绘制气相温度分布
        im3 = axes[1,1].imshow(vapor_temp-273.15, cmap='coolwarm')  # °C
        axes[1,1].set_title('气相温度分布 (°C)')
        plt.colorbar(im3, ax=axes[1,1])
        
        # 保存图形
        plt.tight_layout()
        filename = os.path.join(self.output_dir, f'distribution_{step:06d}.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        
    def _plot_axial_distribution(
        self,
        state: Dict[str, FluidState],
        step: int,
        time: float
    ):
        """
        绘制轴向分布图
        
        参数:
            state: Dict[str, FluidState], 当前状态
            step: int, 当前步数
            time: float, 当前时间
        """
        # 创建图形
        fig, axes = plt.subplots(2, 2, figsize=(12, 12))
        fig.suptitle(f'轴向分布 (t = {time:.3f} s)', fontsize=14)
        
        # 获取轴向位置
        z = np.linspace(0, self.config['grid']['length'], len(state['pressure']))
        
        # 绘制压力分布
        axes[0,0].plot(z, state['pressure']/1e6, 'b-', label='压力')
        axes[0,0].set_xlabel('轴向位置 (m)')
        axes[0,0].set_ylabel('压力 (MPa)')
        axes[0,0].grid(True)
        axes[0,0].legend()
        
        # 绘制空泡率分布
        axes[0,1].plot(z, state['vapor'].alpha, 'r-', label='空泡率')
        axes[0,1].set_xlabel('轴向位置 (m)')
        axes[0,1].set_ylabel('空泡率')
        axes[0,1].grid(True)
        axes[0,1].legend()
        
        # 绘制速度分布
        axes[1,0].plot(z, state['liquid'].u, 'b-', label='液相')
        axes[1,0].plot(z, state['vapor'].u, 'r-', label='气相')
        axes[1,0].set_xlabel('轴向位置 (m)')
        axes[1,0].set_ylabel('速度 (m/s)')
        axes[1,0].grid(True)
        axes[1,0].legend()
        
        # 绘制温度分布
        axes[1,1].plot(z, state['liquid'].T-273.15, 'b-', label='液相')
        axes[1,1].plot(z, state['vapor'].T-273.15, 'r-', label='气相')
        axes[1,1].set_xlabel('轴向位置 (m)')
        axes[1,1].set_ylabel('温度 (°C)')
        axes[1,1].grid(True)
        axes[1,1].legend()
        
        # 保存图形
        plt.tight_layout()
        filename = os.path.join(self.output_dir, f'axial_{step:06d}.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        
    def _plot_time_history(
        self,
        state: Dict[str, FluidState],
        step: int,
        time: float
    ):
        """
        绘制时间历程图
        
        参数:
            state: Dict[str, FluidState], 当前状态
            step: int, 当前步数
            time: float, 当前时间
        """
        # 如果是第一步，初始化时间历程数据
        if step == 0:
            self._initialize_time_history()
            
        # 更新时间历程数据
        self._update_time_history(state, time)
        
        # 创建图形
        fig, axes = plt.subplots(2, 2, figsize=(12, 12))
        fig.suptitle('时间历程', fontsize=14)
        
        # 绘制压力历程
        axes[0,0].plot(self.time_history, self.pressure_history/1e6, 'b-')
        axes[0,0].set_xlabel('时间 (s)')
        axes[0,0].set_ylabel('压力 (MPa)')
        axes[0,0].grid(True)
        
        # 绘制空泡率历程
        axes[0,1].plot(self.time_history, self.void_fraction_history, 'r-')
        axes[0,1].set_xlabel('时间 (s)')
        axes[0,1].set_ylabel('平均空泡率')
        axes[0,1].grid(True)
        
        # 绘制质量流率历程
        axes[1,0].plot(self.time_history, self.mass_flux_history, 'g-')
        axes[1,0].set_xlabel('时间 (s)')
        axes[1,0].set_ylabel('质量流率 (kg/m²s)')
        axes[1,0].grid(True)
        
        # 绘制温度历程
        axes[1,1].plot(self.time_history, self.liquid_temp_history-273.15, 'b-', label='液相')
        axes[1,1].plot(self.time_history, self.vapor_temp_history-273.15, 'r-', label='气相')
        axes[1,1].set_xlabel('时间 (s)')
        axes[1,1].set_ylabel('温度 (°C)')
        axes[1,1].grid(True)
        axes[1,1].legend()
        
        # 保存图形
        plt.tight_layout()
        filename = os.path.join(self.output_dir, f'history_{step:06d}.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        
    def _initialize_time_history(self):
        """初始化时间历程数据"""
        self.time_history = []
        self.pressure_history = []
        self.void_fraction_history = []
        self.mass_flux_history = []
        self.liquid_temp_history = []
        self.vapor_temp_history = []
        
    def _update_time_history(self, state: Dict[str, FluidState], time: float):
        """
        更新时间历程数据
        
        参数:
            state: Dict[str, FluidState], 当前状态
            time: float, 当前时间
        """
        # 添加时间点
        self.time_history.append(time)
        
        # 计算平均值
        self.pressure_history.append(np.mean(state['pressure']))
        self.void_fraction_history.append(np.mean(state['vapor'].alpha))
        
        # 计算质量流率
        mass_flux = (state['liquid'].rho * state['liquid'].u * (1 - state['vapor'].alpha) +
                    state['vapor'].rho * state['vapor'].u * state['vapor'].alpha)
        self.mass_flux_history.append(np.mean(mass_flux))
        
        # 添加温度
        self.liquid_temp_history.append(np.mean(state['liquid'].T))
        self.vapor_temp_history.append(np.mean(state['vapor'].T))
        
    def export_data(self, state: Dict[str, FluidState], step: int, time: float):
        """
        导出数据为CSV格式
        
        参数:
            state: Dict[str, FluidState], 当前状态
            step: int, 当前步数
            time: float, 当前时间
        """
        # 创建数据目录
        data_dir = os.path.join(self.output_dir, 'data')
        os.makedirs(data_dir, exist_ok=True)
        
        # 准备数据
        data = {
            'time': time,
            'step': step,
            'pressure': state['pressure'],
            'void_fraction': state['vapor'].alpha,
            'liquid_velocity': state['liquid'].u,
            'vapor_velocity': state['vapor'].u,
            'liquid_temperature': state['liquid'].T,
            'vapor_temperature': state['vapor'].T
        }
        
        # 保存为CSV文件
        filename = os.path.join(data_dir, f'data_{step:06d}.csv')
        np.savetxt(
            filename,
            np.column_stack([v for v in data.values()]),
            delimiter=',',
            header=','.join(data.keys()),
            comments=''
        ) 